{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"speech_config\": {\n",
    "        \"sample_rate\": 16000,\n",
    "        \"frame_ms\": 25,\n",
    "        \"stride_ms\": 10,\n",
    "        \"num_feature_bins\": 80,\n",
    "        \"feature_type\": \"log_mel_spectrogram\",\n",
    "        \"preemphasis\": 0.97,\n",
    "        \"normalize_signal\": True,\n",
    "        \"normalize_feature\": True,\n",
    "        \"normalize_per_frame\": False,\n",
    "    },\n",
    "    \"decoder_config\": {\n",
    "        \"vocabulary\": None,\n",
    "        \"target_vocab_size\": 1024,\n",
    "        \"max_subword_length\": 4,\n",
    "        \"blank_at_zero\": True,\n",
    "        \"beam_width\": 5,\n",
    "        \"norm_score\": True,\n",
    "    },\n",
    "    \"model_config\": {\n",
    "        \"name\": \"streaming_transducer\",\n",
    "        \"encoder_reductions\": {0: 3, 1: 2},\n",
    "        \"encoder_dmodel\": 320,\n",
    "        \"encoder_rnn_type\": \"lstm\",\n",
    "        \"encoder_rnn_units\": 1024,\n",
    "        \"encoder_nlayers\": 8,\n",
    "        \"encoder_layer_norm\": True,\n",
    "        \"prediction_embed_dim\": 320,\n",
    "        \"prediction_embed_dropout\": 0.0,\n",
    "        \"prediction_num_rnns\": 2,\n",
    "        \"prediction_rnn_units\": 1024,\n",
    "        \"prediction_rnn_type\": \"lstm\",\n",
    "        \"prediction_projection_units\": 320,\n",
    "        \"prediction_layer_norm\": True,\n",
    "        \"joint_dim\": 320,\n",
    "        \"joint_activation\": \"tanh\",\n",
    "    },\n",
    "    \"learning_config\": {\n",
    "        \"train_dataset_config\": {\n",
    "            \"use_tf\": True,\n",
    "            \"augmentation_config\": {\n",
    "                \"feature_augment\": {\n",
    "                    \"time_masking\": {\n",
    "                        \"num_masks\": 10,\n",
    "                        \"mask_factor\": 100,\n",
    "                        \"p_upperbound\": 0.05,\n",
    "                    },\n",
    "                    \"freq_masking\": {\"num_masks\": 1, \"mask_factor\": 27},\n",
    "                }\n",
    "            },\n",
    "            \"data_paths\": [\n",
    "                \"/mnt/h/ML/Datasets/ASR/Raw/LibriSpeech/train-clean-100/transcripts.tsv\"\n",
    "            ],\n",
    "            \"tfrecords_dir\": None,\n",
    "            \"shuffle\": True,\n",
    "            \"cache\": True,\n",
    "            \"buffer_size\": 100,\n",
    "            \"drop_remainder\": True,\n",
    "            \"stage\": \"train\",\n",
    "        },\n",
    "        \"eval_dataset_config\": {\n",
    "            \"use_tf\": True,\n",
    "            \"data_paths\": None,\n",
    "            \"tfrecords_dir\": None,\n",
    "            \"shuffle\": False,\n",
    "            \"cache\": True,\n",
    "            \"buffer_size\": 100,\n",
    "            \"drop_remainder\": True,\n",
    "            \"stage\": \"eval\",\n",
    "        },\n",
    "        \"test_dataset_config\": {\n",
    "            \"use_tf\": True,\n",
    "            \"data_paths\": [\n",
    "                \"/mnt/h/ML/Datasets/ASR/Raw/LibriSpeech/test-clean/transcripts.tsv\"\n",
    "            ],\n",
    "            \"tfrecords_dir\": None,\n",
    "            \"shuffle\": False,\n",
    "            \"cache\": True,\n",
    "            \"buffer_size\": 100,\n",
    "            \"drop_remainder\": True,\n",
    "            \"stage\": \"test\",\n",
    "        },\n",
    "        \"optimizer_config\": {\"class_name\": \"adam\", \"config\": {\"learning_rate\": 0.0001}},\n",
    "        \"running_config\": {\n",
    "            \"batch_size\": 2,\n",
    "            \"num_epochs\": 20,\n",
    "            \"checkpoint\": {\n",
    "                \"filepath\": \"/mnt/e/Models/local/rnn_transducer/checkpoints/{epoch:02d}.h5\",\n",
    "                \"save_best_only\": False,\n",
    "                \"save_weights_only\": True,\n",
    "                \"save_freq\": \"epoch\",\n",
    "            },\n",
    "            \"states_dir\": \"/mnt/e/Models/local/rnn_transducer/states\",\n",
    "            \"tensorboard\": {\n",
    "                \"log_dir\": \"/mnt/e/Models/local/rnn_transducer/tensorboard\",\n",
    "                \"histogram_freq\": 1,\n",
    "                \"write_graph\": True,\n",
    "                \"write_images\": True,\n",
    "                \"update_freq\": \"epoch\",\n",
    "                \"profile_batch\": 2,\n",
    "            },\n",
    "        },\n",
    "    },\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata = {\n",
    "    \"train\": {\"max_input_length\": 2974, \"max_label_length\": 194, \"num_entries\": 281241},\n",
    "    \"eval\": {\"max_input_length\": 3516, \"max_label_length\": 186, \"num_entries\": 5567},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import argparse\n",
    "from tensorflow_asr.utils import env_util\n",
    "\n",
    "env_util.setup_environment()\n",
    "import tensorflow as tf\n",
    "\n",
    "tf.keras.backend.clear_session()\n",
    "tf.config.optimizer.set_experimental_options({\"auto_mixed_precision\": True})\n",
    "strategy = env_util.setup_strategy([0])\n",
    "\n",
    "from tensorflow_asr.configs.config import Config\n",
    "from tensorflow_asr.datasets import asr_dataset\n",
    "from tensorflow_asr.featurizers import speech_featurizers, text_featurizers\n",
    "from tensorflow_asr.models.transducer.rnn_transducer import RnnTransducer\n",
    "from tensorflow_asr.optimizers.schedules import TransformerSchedule\n",
    "\n",
    "config = Config(config)\n",
    "speech_featurizer = speech_featurizers.TFSpeechFeaturizer(config.speech_config)\n",
    "\n",
    "text_featurizer = text_featurizers.CharFeaturizer(config.decoder_config)\n",
    "\n",
    "train_dataset = asr_dataset.ASRSliceDataset(\n",
    "    speech_featurizer=speech_featurizer,\n",
    "    text_featurizer=text_featurizer,\n",
    "    **vars(config.learning_config.train_dataset_config),\n",
    "    indefinite=True\n",
    ")\n",
    "eval_dataset = asr_dataset.ASRSliceDataset(\n",
    "    speech_featurizer=speech_featurizer,\n",
    "    text_featurizer=text_featurizer,\n",
    "    **vars(config.learning_config.eval_dataset_config),\n",
    "    indefinite=True\n",
    ")\n",
    "\n",
    "train_dataset.load_metadata(metadata)\n",
    "eval_dataset.load_metadata(metadata)\n",
    "speech_featurizer.reset_length()\n",
    "text_featurizer.reset_length()\n",
    "\n",
    "global_batch_size = config.learning_config.running_config.batch_size\n",
    "global_batch_size *= strategy.num_replicas_in_sync\n",
    "\n",
    "train_data_loader = train_dataset.create(global_batch_size)\n",
    "eval_data_loader = eval_dataset.create(global_batch_size)\n",
    "\n",
    "with strategy.scope():\n",
    "    # build model\n",
    "    rnnt = RnnTransducer(**config.model_config, vocabulary_size=text_featurizer.num_classes)\n",
    "    rnnt.make(speech_featurizer.shape)\n",
    "    rnnt.summary(line_length=100)\n",
    "\n",
    "    rnnt.compile(\n",
    "        optimizer=config.learning_config.optimizer_config,\n",
    "        experimental_steps_per_execution=10,\n",
    "        global_batch_size=global_batch_size,\n",
    "        blank=text_featurizer.blank\n",
    "    )\n",
    "\n",
    "callbacks = [\n",
    "    tf.keras.callbacks.ModelCheckpoint(**config.learning_config.running_config.checkpoint),\n",
    "    tf.keras.callbacks.experimental.BackupAndRestore(config.learning_config.running_config.states_dir),\n",
    "    tf.keras.callbacks.TensorBoard(**config.learning_config.running_config.tensorboard)\n",
    "]\n",
    "\n",
    "rnnt.fit(\n",
    "    train_data_loader,\n",
    "    epochs=config.learning_config.running_config.num_epochs,\n",
    "    validation_data=eval_data_loader,\n",
    "    callbacks=callbacks,\n",
    "    steps_per_epoch=train_dataset.total_steps,\n",
    "    validation_steps=eval_dataset.total_steps\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.8 64-bit ('tfo': conda)",
   "name": "python388jvsc74a57bd045f983f364f7a4cc7101e6d6987a2125bf0c2b5c5c9855ff35103689f542d13f"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8-final"
  },
  "metadata": {
   "interpreter": {
    "hash": "45f983f364f7a4cc7101e6d6987a2125bf0c2b5c5c9855ff35103689f542d13f"
   }
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
